// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#pragma once

#include <pollux/dwio/common/data_buffer_holder.h>
#include <pollux/dwio/common/wrap/zero-copy-stream-wrapper.h>

namespace kumo::pollux::dwio::common {
    /**
     * Record write position for creating index stream
     */
    class PositionRecorder {
    public:
        virtual ~PositionRecorder() = default;

        virtual void add(uint64_t pos, int32_t entry) = 0;
    };

    /**
     * A subclass of Google's ZeroCopyOutputStream that supports output to memory
     * buffer, and flushing to OutputStream.
     * By extending Google's class, we get the ability to pass it directly
     * to the protobuf writers.
     */
    class BufferedOutputStream : public google::protobuf::io::ZeroCopyOutputStream {
    public:
        explicit BufferedOutputStream(DataBufferHolder &bufferHolder)
            : bufferHolder_{bufferHolder}, buffer_{bufferHolder_.getMemoryPool()} {
        }

        bool Next(void **data, int32_t *size) override {
            return Next(data, size, 1);
        }

        void BackUp(int32_t count) override;

        google::protobuf::int64 ByteCount() const override {
            return static_cast<google::protobuf::int64>(size());
        }

        bool WriteAliasedRaw(const void * /* unused */, int32_t /* unused */)
        override {
            POLLUX_NYI("WriteAliasedRaw is not supported");
            return false;
        }

        bool AllowsAliasing() const override {
            return false;
        }

        /// Similar to Next(data, size), but in addition, allows callers to specify
        /// expected increment so buffer allocation is more efficient
        virtual bool Next(void **data, int32_t *size, uint64_t increment);

        virtual uint64_t flush();

        virtual std::string getName() const {
            return melon::to<std::string>("BufferedOutputStream, size: ", size());
        }

        virtual uint64_t size() const {
            // buffered size + flushed size
            return buffer_.size() + bufferHolder_.size();
        }

        virtual void recordPosition(
            PositionRecorder &recorder,
            int32_t bufferLength,
            int32_t bufferOffset,
            int32_t strideIndex = -1) const {
            auto streamSize = size();
            if (streamSize > 0) {
                streamSize -= (bufferLength - bufferOffset);
            }
            recorder.add(streamSize, strideIndex);
        }

    protected:
        // try increase buffer size, and then assign to output buffer/size. Returns
        // false if buffer size remained the same
        bool tryResize(
            void **buffer,
            int32_t *size,
            uint64_t headerSize,
            uint64_t increment) {
            auto origSize = buffer_.size();
            if (bufferHolder_.tryResize(buffer_, headerSize, increment)) {
                // if original buffer is empty. need to adjust buffer position/size to
                // accommodate header
                if (MELON_UNLIKELY(!origSize && headerSize)) {
                    origSize = headerSize;
                }
                *buffer = buffer_.data() + origSize;
                *size = static_cast<int32_t>(buffer_.size() - origSize);
                return true;
            } else {
                return false;
            }
        }

        // this method is called after tryResize() returns false
        void flushAndReset(
            void **buffer,
            int32_t *size,
            uint64_t headerSize,
            const std::vector<melon::StringPiece> &bufferToFlush) {
            bufferHolder_.take(bufferToFlush);
            *buffer = buffer_.data() + headerSize;
            *size = static_cast<int32_t>(buffer_.size() - headerSize);
        }

        DataBufferHolder &bufferHolder_;
        dwio::common::DataBuffer<char> buffer_;
    };

    /**
     * An append only buffered stream that allows buffer, and flushing to
     * OutputStream. By extending Google's class, we get the ability to pass it
     * directly to the protobuf writers.
     */
    class AppendOnlyBufferedStream {
    public:
        explicit AppendOnlyBufferedStream(
            std::unique_ptr<BufferedOutputStream> outStream)
            : outStream_(std::move(outStream)) {
            POLLUX_CHECK_NOT_NULL(outStream_);
        }

        void write(const char *data, size_t size);

        uint64_t flush();

        uint64_t size() const {
            // size of written content, not size of allocated. So need to subtract size
            // that is not in use
            return outStream_->size() - (bufferLength_ - bufferOffset_);
        }

        void recordPosition(PositionRecorder &recorder, int32_t strideIndex = -1)
        const {
            outStream_->recordPosition(
                recorder, bufferLength_, bufferOffset_, strideIndex);
        }

    private:
        std::unique_ptr<BufferedOutputStream> outStream_;
        char *buffer_{nullptr};
        int32_t bufferLength_{0};
        int32_t bufferOffset_{0};
    };
} // namespace kumo::pollux::dwio::common
